/*
 * Copyright (c) Kumo Inc. and affiliates.
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <melon/futures/future.h>

#include <melon/likely.h>
#include <melon/singleton.h>
#include <melon/futures/heap_timekeeper.h>
#include <melon/futures/thread_wheel_timekeeper.h>
#include <melon/config.h>

namespace melon {
    namespace futures {
        namespace detail {
            void WithinInterruptHandler::operator()(exception_wrapper const &ew) const {
                if (auto locked = ptr.lock()) {
                    locked->raise(ew);
                }
            }

            void WithinAfterFutureCallback::operator()(Try<Unit> &&t) {
                if (t.hasException() &&
                    t.exception().is_compatible_with<FutureCancellation>()) {
                    // This got cancelled by thisFuture so we can just return.
                    return;
                }

                auto lockedCtx = ctx.lock();
                if (!lockedCtx) {
                    // ctx already released. "this" completed first, cancel "after"
                    return;
                }
                // "after" completed first, cancel "this"
                lockedCtx->thisFuture.raise(FutureTimeout());
                if (!lockedCtx->token.exchange(true, std::memory_order_relaxed)) {
                    auto &exn = t.hasException() ? t.exception() : lockedCtx->exception;
                    lockedCtx->doPromiseSetException(*lockedCtx, std::move(exn));
                }
            }
        } // namespace detail

        SemiFuture<Unit> sleep(HighResDuration dur, Timekeeper *tk) {
            std::shared_ptr<Timekeeper> tks;
            if (MELON_LIKELY(!tk)) {
                tks = melon::detail::getTimekeeperSingleton();
                tk = tks.get();
            }

            if (MELON_UNLIKELY(!tk)) {
                return makeSemiFuture<Unit>(FutureNoTimekeeper());
            }

            return tk->after(dur);
        }

        Future<Unit> sleepUnsafe(HighResDuration dur, Timekeeper *tk) {
            return sleep(dur, tk).toUnsafeFuture();
        }

        namespace {
            template<typename Ptr>
            class FutureWaiter : public fibers::Baton::Waiter {
            public:
                FutureWaiter(Promise<Unit> promise, Ptr baton)
                    : promise_(std::move(promise)), baton_(std::move(baton)) {
                    baton_->setWaiter(*this);
                }

                void post() override {
                    promise_.setValue();
                    delete this;
                }

            private:
                Promise<Unit> promise_;
                Ptr baton_;
            };
        } // namespace

        SemiFuture<Unit> wait(std::unique_ptr<fibers::Baton> baton) {
            Promise<Unit> promise;
            auto sf = promise.getSemiFuture();
            new FutureWaiter<std::unique_ptr<fibers::Baton> >(
                std::move(promise), std::move(baton));
            return sf;
        }

        SemiFuture<Unit> wait(std::shared_ptr<fibers::Baton> baton) {
            Promise<Unit> promise;
            auto sf = promise.getSemiFuture();
            new FutureWaiter<std::shared_ptr<fibers::Baton> >(
                std::move(promise), std::move(baton));
            return sf;
        }
    } // namespace futures

    namespace detail {
        namespace {
            Singleton<Timekeeper, TimekeeperSingletonTag> gTimekeeperSingleton(
                []() -> Timekeeper * {
                    if (turbo::get_flag(FLAGS_melon_futures_use_thread_wheel_timekeeper)) {
                        return new ThreadWheelTimekeeper;
                    } else {
                        return new HeapTimekeeper;
                    }
                });
        } // namespace

        std::shared_ptr<Timekeeper> getTimekeeperSingleton() {
            return gTimekeeperSingleton.try_get();
        }
    } // namespace detail

#if MELON_USE_EXTERN_FUTURE_UNIT
namespace futures {
namespace detail {
template class FutureBase<Unit>;
} // namespace detail
} // namespace futures

template class Future<Unit>;
template class SemiFuture<Unit>;
#endif
} // namespace melon
